package org.ragdollcat.secondaiagent.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.ragdollcat.secondaiagent.agent.enums.AgentState;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.List;
import java.util.stream.Collectors;

@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent {

    /**
     * 定义可用的工具
     */
    private ToolCallback[] availableTools;

    /**
     * 定义可用的工具，这里使用mcp
      */
//    private ToolCallbackProvider toolCallbackProvider;

    /**
     * 定义保存工具调用信息的响应结果
     */
    private ChatResponse toolCallChatResponse;


    /**
     * 定义工具调用的管理者
     */
    private ToolCallingManager toolCallingManager;


    /**
     * 禁用框架自动调用工具的上下文
     */
    private ChatOptions chatOptions;


    public ToolCallAgent(ToolCallback[] availableTools) {
        super();
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        //自己维护上下文
        this.chatOptions = DashScopeChatOptions.builder().withProxyToolCalls(true).build();
    }

    /**
     * 思考 判断是否需要调用工具
     *
     * @return
     */
    @Override
    public boolean think(SseEmitter sseEmitter) throws IOException {
        //下一步提示词不为空
        if (StrUtil.isNotBlank(getNextStepPrompt())) {
            //作为用户提示词，加入到上下文记忆中
            getMemoryMessages().add(new UserMessage(getNextStepPrompt()));
        }
        //得到最新的上下文
        List<Message> memoryMessages = getMemoryMessages();

        Prompt prompt = new Prompt(memoryMessages, chatOptions);
        try {
            //获取带有工具选项的响应
            toolCallChatResponse = getChatClient().
                    //预设prompt
                            prompt(prompt).
                    //系统提示词
                            system(getSystemPrompt())
                    //设置工具
                    .tools(availableTools)
                    .call()
                    .chatResponse();

            //得到助手提示词
            AssistantMessage assistantMessage = toolCallChatResponse.getResult().getOutput();

            String text = assistantMessage.getText();
            List<AssistantMessage.ToolCall> toolCalls = assistantMessage.getToolCalls();
            log.info("{}的思考:{}", getName(), text);
            log.info("{}选择了{}个工具使用", getName(), toolCalls.size());
            String toolInfo = toolCalls.stream().map(toolCall -> "工具名称" + toolCall.name() + "工具参数" + toolCall.arguments()).collect(Collectors.joining("/n"));
            log.info(toolInfo);

            if (CollUtil.isEmpty(toolCalls)) {
                log.info("{}本次无需使用工具", getName());
                getMemoryMessages().add(assistantMessage);
                sseEmitter.send(String.format("%s的思考:%s", getName(),text));
                return false;
            }
            sseEmitter.send(String.format("%s的思考:%s", getName(),text));
            return true;
        } catch (Exception e) {
            log.error("{}的思考过程遇到了问题：",getName(),e);
            //记录上下文
            getMemoryMessages().add(new AssistantMessage("处理时遇到错误" + e.getMessage()));
            sseEmitter.send("处理时遇到错误" + e.getMessage());
            return false;
        }
    }

    /**
     * 行动 调用工具
     *
     * @return
     */
    @Override
    public String act() {
        if (!toolCallChatResponse.hasToolCalls()){
            return "没有工具调用";
        }
        //准备调用工具
        Prompt prompt = new Prompt(getMemoryMessages(), chatOptions);
        //调用工具
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);

        //记录消息上下文 conversationHistory中已经包含了历史的消息
        setMemoryMessages(toolExecutionResult.conversationHistory());

        //返回当前工具调用的结果
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        String resp = toolResponseMessage.getResponses().stream().map(toolResponse -> "工具" + toolResponse.name() + "完成了工作，结果" + toolResponse.responseData()).collect(Collectors.joining());

        //判断是否调用了终止工具
        boolean aniedMatch = toolResponseMessage.getResponses().stream().anyMatch(toolResponse -> "doTerminate".equals(toolResponse.name()));
        if (aniedMatch){
            //修改状态为终止
            setAgentState(AgentState.FINISHED);
        }


        log.info(resp);

        return resp;
    }
}
